import argparse  # 导入 argparse 模块，用于解析命令行参数
import numpy as np  # 导入 NumPy 库，用于处理数组数据
from pathlib import Path  # 导入 Path 模块，用于处理文件路径
import torch  # 导入 PyTorch 库，用于深度学习模型推理
from sklearn.metrics import classification_report, ConfusionMatrixDisplay  # 导入分类报告和混淆矩阵显示工具
import matplotlib.pyplot as plt  # 导入 Matplotlib 库，用于绘图
from crnn import WakeWordCRNN  # 导入自定义的 CRNN 模型（WakeWordCRNN）
from dataset import WakeWordDataset  # 导入自定义的数据集类（WakeWordDataset）


def evaluate(args):  # 定义评估函数，接收命令行参数 args
    device = torch.device(args.device)  # 根据传入的参数选择设备（CPU 或 GPU）

    # 加载模型
    model = WakeWordCRNN().to(device)  # 初始化 CRNN 模型，并将其移动到指定设备
    model.load_state_dict(torch.load(args.model_path, map_location=device))  # 加载预训练模型权重
    model.eval()  # 设置模型为评估模式

    # 加载数据集
    test_dataset = WakeWordDataset(args.test_dir, augment=False)  # 加载测试数据集，不启用数据增强
    test_loader = test_dataset.get_loader(args.batch_size, shuffle=False)  # 创建测试数据加载器，关闭打乱顺序

    # 推理
    all_labels = []  # 初始化存储真实标签的列表
    all_preds = []  # 初始化存储预测标签的列表
    all_paths = []  # 初始化存储样本路径的列表

    with torch.no_grad():  # 禁用梯度计算，节省内存
        for features, labels, paths in test_loader:  # 遍历测试数据加载器
            features = features.to(device)  # 将特征移动到指定设备
            outputs = model(features)  # 前向传播，计算模型输出
            _, preds = torch.max(outputs, 1)  # 获取预测结果（取最大值索引）

            all_labels.extend(labels.cpu().numpy())  # 将真实标签添加到列表中
            all_preds.extend(preds.cpu().numpy())  # 将预测标签添加到列表中
            all_paths.extend(paths)  # 将样本路径添加到列表中

    # 生成报告
    print(classification_report(all_labels, all_preds, target_names=['not_wake', 'wake']))  # 打印分类报告

    # 绘制混淆矩阵
    fig, ax = plt.subplots(figsize=(6, 6))  # 创建绘图区域，设置图像大小
    ConfusionMatrixDisplay.from_predictions(
        all_labels, all_preds,  # 使用真实标签和预测标签绘制混淆矩阵
        display_labels=['Not Wake', 'Wake'],  # 设置类别标签
        cmap='Blues', ax=ax  # 设置颜色映射和绘图轴
    )
    plt.savefig(Path(args.save_dir) / "confusion_matrix.png")  # 保存混淆矩阵图像

    # 保存错误样本
    with open(Path(args.save_dir) / "errors.txt", 'w') as f:  # 打开文件，准备写入错误样本信息
        for path, true, pred in zip(all_paths, all_labels, all_preds):  # 遍历所有样本路径、真实标签和预测标签
            if true != pred:  # 如果真实标签与预测标签不一致
                f.write(f"{path}\tTrue: {true}\tPred: {pred}\n")  # 写入错误样本的路径和标签信息


if __name__ == "__main__":  # 主程序入口
    parser = argparse.ArgumentParser()  # 创建命令行参数解析器
    parser.add_argument("--test_dir", default="dataset/train", required=True)  # 添加测试数据目录参数
    parser.add_argument("--model_path", default="checkpoints/best_model.pth", required=True)  # 添加模型路径参数
    parser.add_argument("--batch_size", type=int, default=50)  # 添加批量大小参数
    parser.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")  # 添加设备参数
    parser.add_argument("--save_dir", default="./results")  # 添加结果保存目录参数
    args = parser.parse_args()  # 解析命令行参数

    Path(args.save_dir).mkdir(exist_ok=True)  # 确保结果保存目录存在
    evaluate(args)  # 调用评估函数，传入解析后的参数
    #这段代码实现了一个完整的模型评估流程，包括加载预训练模型、推理测试数据集、生成分类报告、绘制混淆矩阵以及保存错误样本。